from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.ticker import FuncFormatter
import mplcursors

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

# 全局变量，用于存储和图像显示线程交互的数据
class imgview:
    global_image = None
    imgview_exit = False
    trajectory = []
    focus_i = 0
    traj_i = 0

imgview_data = imgview()

# 定义一个函数，用于在独立线程中显示图像
def show_image():
    grid_size_display = 20
    while not imgview_data.imgview_exit:
        # 检查全局变量是否有图像
        if imgview_data.global_image is not None:
            img = np.copy(imgview_data.global_image)
            state_x = imgview_data.trajectory[imgview_data.traj_i][0]
            state_y = imgview_data.trajectory[imgview_data.traj_i][1]
            cv2.circle(img, (state_y * grid_size_display + int(grid_size_display/2), state_x * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
            # 显示图像
            cv2.imshow("Image", img)
            key = cv2.waitKey(1)
            if key == ord('a'):
                imgview_data.focus_i -= 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
            elif key == ord('d'):
                imgview_data.focus_i += 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
        else:
            # 图像还未产生，等待100毫秒
            time.sleep(0.1)

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)
    parser.add_argument("--load_data", type=int, default=rpl_config.load_data)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i
    rpl_config.load_data = args.load_data

    if rpl_config.load_data == 0:

        """ load model
        """
        params = load_weights(rpl_config.model_pth)
        # 定义一个函数，用于生成随机权重
        def init_weights_r(key, shape):
            return jax.random.normal(key, shape)
        # 生成一个随机的 PRNGKey
        key = jax.random.PRNGKey(np.random.randint(0, 1000))
        # key = jax.random.PRNGKey(4512)
        # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
        random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
        # params = random_params

        # get elements of params
        tree_leaves = jax.tree_util.tree_leaves(params)
        for i in range(len(tree_leaves)):
            print("shape of leaf ", i, ": ", tree_leaves[i].shape)

        bias1 = np.array(tree_leaves[0])
        mat1 = np.array(tree_leaves[1])
        print("mat1.shape: ", mat1.shape)
        print("bias1: ", bias1)

        mat_obs = np.array(tree_leaves[1])[rpl_config.nn_size:rpl_config.nn_size+10,:]
        mat_intr = np.array(tree_leaves[1])[0:rpl_config.nn_size,:]
        print("mat_obs.shape: ", mat_obs.shape)
        
        """ create agent
        """
        if rpl_config.nn_type == "vanilla":
            model = RNN(hidden_dims = rpl_config.nn_size)
        elif rpl_config.nn_type == "gru":
            model = GRU(hidden_dims = rpl_config.nn_size)

        # check if param fits the agent
        if rpl_config.nn_type == "vanilla":
            assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

        n_samples = 10000
        k1 = npr.randint(0, 1000000)
        rnn_state = model.initial_state_rnd(n_samples, k1)
        rnn_state_old = rnn_state.copy()
        diff = jnp.abs(rnn_state - rnn_state_old)
        rnn_state_old = rnn_state.copy()
        diff_norm = jnp.linalg.norm(diff, axis=1)
        diff_norm_old = diff_norm.copy()
        norm_std = diff_norm.copy()

        rnn_state_init = rnn_state.copy()

        obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

        rnn_state_trajectory = []
        intr_field = []

        for t in range(rpl_config.life_duration):

            progress_bar(t, rpl_config.life_duration)

            intr_vector = np.dot(rnn_state, mat_intr) + bias1
            intr_vector = np.tanh(intr_vector)
            intr_field.append(intr_vector - rnn_state)

            rnn_state_trajectory.append(np.array(rnn_state).copy())
            
            """ model forward 
            """
            # rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
            rnn_state = intr_vector
            
        print(rnn_state.shape)
        print(norm_std.shape)
        rnn_state_np = np.array(rnn_state)

        # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
        rnn_state_trajectory_np = np.array(rnn_state_trajectory)
        rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
        print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

        # 将 intr_field 展开成 rnn_state_np 的形状
        intr_field_np = np.array(intr_field)
        intr_field_np = intr_field_np.reshape(-1, intr_field_np.shape[-1])
        print("shape of intr_field_np: ", intr_field_np.shape)

        # 对 rnn_state_np 进行 PCA
        intrinsic_pca = PCA()
        # pca.fit(rnn_state_np)
        intrinsic_pca.fit(rnn_state_trajectory_np)

        # # 打印 variance ratio
        # print(pca.explained_variance_ratio_)

        rnn_state_trajectory_np_pca = intrinsic_pca.transform(rnn_state_trajectory_np)

        print("shape of rnn_state_trajectory_np_pca: ", rnn_state_trajectory_np_pca.shape)

        @jax.jit
        def compute_intr_field(HS_trajectory_1, mat_intr, bias1):
            intr_vector = jnp.dot(HS_trajectory_1, mat_intr) + bias1
            intr_vector = jnp.tanh(intr_vector)
            intr_field = intr_vector - HS_trajectory_1
            return intr_field
        compute_intr_field_vmap = vmap(compute_intr_field, in_axes=(0, None, None))

        contributions_of_incremental_field = []

        # 1. 计算宿主空间向量场
        compute_dims = [0]

        # 将除了指定维度以外的所有维度置零
        mask = np.zeros(rnn_state_trajectory_np_pca.shape[1], dtype=bool)
        mask[compute_dims] = True
        rnn_state_trajectory_np_pca_masked = rnn_state_trajectory_np_pca.copy()
        rnn_state_trajectory_np_pca_masked[:, ~mask] = 0

        # 将 pca 空间的点，逆变换回原始相空间
        HS_trajectory_1 = intrinsic_pca.inverse_transform(rnn_state_trajectory_np_pca_masked)

        # 计算主成分空间中的内禀向量场
        intr_field1 = compute_intr_field_vmap(HS_trajectory_1, mat_intr, bias1)

        @jax.jit
        def comupte_contribution(_vector, _base):
            return jnp.dot(_vector, _base) / jnp.linalg.norm(_base)
        
        compute_contribution_vmap = vmap(comupte_contribution, in_axes=(0, 0))

        contributions = compute_contribution_vmap(intr_field1, intr_field_np)

        # print("shape of contributions: ", contributions.shape)

        contributions_of_incremental_field.append(np.array(contributions))


        for i in range(1, rpl_config.nn_size):

            progress_bar(i, rpl_config.nn_size)

            # 2. 计算加入扰动维度后的宿主空间的向量场
            compute_dims = [d for d in range(i+1)]

            # 将除了指定维度以外的所有维度置零
            mask = np.zeros(rnn_state_trajectory_np_pca.shape[1], dtype=bool)
            mask[compute_dims] = True
            rnn_state_trajectory_np_pca_masked = rnn_state_trajectory_np_pca.copy()
            rnn_state_trajectory_np_pca_masked[:, ~mask] = 0

            # 将 pca 空间的点，逆变换回原始相空间
            HS_trajectory_2 = intrinsic_pca.inverse_transform(rnn_state_trajectory_np_pca_masked)

            # 计算主成分空间中的内禀向量场
            intr_field2 = compute_intr_field_vmap(HS_trajectory_2, mat_intr, bias1)

            # 计算 IVF2 和 IVF1 的中间向量
            incremental_field = intr_field2 - intr_field1

            # 量化中间向量的贡献度
            contributions = compute_contribution_vmap(incremental_field, intr_field_np)
            contributions_of_incremental_field.append(np.array(contributions))

            intr_field1 = intr_field2.copy()

        contributions_of_incremental_field = np.array(contributions_of_incremental_field)
        print("shape of contributions_of_incremental_field: ", contributions_of_incremental_field.shape)
        print("memory usage: ", sys.getsizeof(contributions_of_incremental_field)/1024/1024, "MB")

        # 交换 contributions_of_incremental_field 的第一、第二维
        contributions_of_incremental_field = np.swapaxes(contributions_of_incremental_field, 0, 1)
        print("shape of contributions_of_incremental_field: ", contributions_of_incremental_field.shape)

        max_sum = jnp.sum(jnp.maximum(contributions_of_incremental_field, 0), axis=1)
        print("shape of max_sum: ", max_sum.shape)

        # 归一化每个向量
        normalized_contributions = contributions_of_incremental_field / max_sum[:, jnp.newaxis]

        # 保存数据
        np.save("./logs/rnn_state_trajectory_np_pca.npy", rnn_state_trajectory_np_pca)
        np.save("./logs/normalized_contributions.npy", normalized_contributions)

    elif rpl_config.load_data == 1:
        rnn_state_trajectory_np_pca = np.load("./logs/rnn_state_trajectory_np_pca.npy")
        normalized_contributions = np.load("./logs/normalized_contributions.npy")

    view_dim0 = 0
    view_dim1 = 1
    view_dim2 = 2
    
    rnn_state_trajectory_np_pca_subset = []
    normalized_contributions_subset = []
    
    for i in range(rnn_state_trajectory_np_pca.shape[0] - 1):
        if i % 500 == 0:
            rnn_state_trajectory_np_pca_subset.append(rnn_state_trajectory_np_pca[i])
            normalized_contributions_subset.append(normalized_contributions[i,:])

    normalized_contributions_subset = np.array(normalized_contributions_subset)
    # 交换 info_dim_ratio_subset 的第一个和第二个维度
    normalized_contributions_subset = np.swapaxes(normalized_contributions_subset, 0, 1)
    print("shape of info_dim_ratio_subset: ", normalized_contributions_subset.shape)

    for info_dim_ in range(rpl_config.nn_size//4):

        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, subplot_kw={'projection': '3d'})

        info_dim = info_dim_*4+0
        scatter1 = ax1.scatter(np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim0], 
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim1], 
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim2], 
                c=normalized_contributions_subset[info_dim], cmap='seismic'
                , alpha=0.5
                )
        fig.colorbar(scatter1)
        cursor1 = mplcursors.cursor(ax1, hover=True)
        @cursor1.connect("add")
        def on_add(sel):
            index = sel.target.index
            x, y, z = sel.artist._offsets3d
            sel.annotation.set_alpha(0.9)
            sel.annotation.set_text(f"({x[index]:.2f}, {y[index]:.2f}, {z[index]:.2f})"+" IVF contribution: "+
                                    str(int(normalized_contributions_subset[info_dim][index]*100))+"%")
        
        info_dim1 = info_dim_*4+1
        scatter2 = ax2.scatter(np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim0],
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim1],
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim2],
                c=normalized_contributions_subset[info_dim1], cmap='seismic'
                , alpha=0.5
                )
        fig.colorbar(scatter2)
        cursor2 = mplcursors.cursor(ax2, hover=True)
        @cursor2.connect("add")
        def on_add(sel):
            index = sel.target.index
            x, y, z = sel.artist._offsets3d
            sel.annotation.set_alpha(0.9)
            sel.annotation.set_text(f"({x[index]:.2f}, {y[index]:.2f}, {z[index]:.2f})"+" IVF contribution: "+
                                    str(int(normalized_contributions_subset[info_dim1][index]*100))+"%")
            
        info_dim2 = info_dim_*4+2
        scatter3 = ax3.scatter(np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim0],
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim1],
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim2],
                c=normalized_contributions_subset[info_dim2], cmap='seismic'
                , alpha=0.5
                )
        fig.colorbar(scatter3)
        cursor3 = mplcursors.cursor(ax3, hover=True)
        @cursor3.connect("add")
        def on_add(sel):
            index = sel.target.index
            x, y, z = sel.artist._offsets3d
            sel.annotation.set_alpha(0.9)
            sel.annotation.set_text(f"({x[index]:.2f}, {y[index]:.2f}, {z[index]:.2f})"+" IVF contribution: "+
                                    str(int(normalized_contributions_subset[info_dim2][index]*100))+"%")
            
        info_dim3 = info_dim_*4+3
        scatter4 = ax4.scatter(np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim0],
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim1],
                np.array(rnn_state_trajectory_np_pca_subset)[:, view_dim2],
                c=normalized_contributions_subset[info_dim3], cmap='seismic'
                , alpha=0.5
                )
        fig.colorbar(scatter4)
        cursor4 = mplcursors.cursor(ax4, hover=True)
        @cursor4.connect("add")
        def on_add(sel):
            index = sel.target.index
            x, y, z = sel.artist._offsets3d
            sel.annotation.set_alpha(0.9)
            sel.annotation.set_text(f"({x[index]:.2f}, {y[index]:.2f}, {z[index]:.2f})"+" IVF contribution: "+
                                    str(int(normalized_contributions_subset[info_dim3][index]*100))+"%")
            
        # 为 ax1, ax2, ax3, ax4 设置标题
        ax1.set_title("contribution of PC: " + str(info_dim))
        ax2.set_title("contribution of PC: " + str(info_dim1))
        ax3.set_title("contribution of PC: " + str(info_dim2))
        ax4.set_title("contribution of PC: " + str(info_dim3))

        plt.show()

if __name__ == "__main__":
    
    ivf()